Motivation

Torch-TRT compiles PyTorch models to TensorRT engines, but today there is no first-class path for users who want to replace a subgraph with their own Triton, CuTile, or CuTeDSL kernel inside the compiled engine. The typical workaround — writing a C++ TRT plugin and registering it manually — requires leaving the Python ecosystem, managing separate build systems, and wiring up the plugin registry by hand. This is a significant barrier for researchers and ML engineers who already have high-performance Python kernels.

TensorRT 10.x introduced Quick Deployable Plugins (QDP), which support AOT-compiled Python kernels (@trtp.aot_impl) that are embedded directly into the TRT engine with no Python required at runtime. This PR adds the descriptor and registration layer that lets users express a custom QDP plugin as a plain Python object and pass it to Torch-TRT.

What's in this PR

Public API (import torch_tensorrt.annotation as tta):

QDP registration (_custom_plugin/):

• _descriptor.py — CustomPluginSpec dataclass + custom_plugin() factory; computes a deterministic op name from the kernel function identity + config hash; register_custom_plugin() registers @trtp.register / @trtp.autotune / @trtp.aot_impl with TRT's process-global QDP registry using double-checked locking for xdist safety.
• _lowering.py — lowers a CustomPluginSpec to a TRT plugin layer via ctx.net.add_plugin(trtp.op.<ns>.<name>(*inputs), aot=True); injects weight tensors as add_constant layers.
• _qdp_utils.py — deterministic op-name derivation, tactic table building, meta-tensor helpers for symbolic shape inference.
• _symbolic.py — SymbolicTensor abstraction for QDP shape/dtype descriptor registration.

AOT backends (_custom_plugin/_aot/):

• _triton.py — Triton → PTX via triton.compile; per-config tactic entries.
• _cutile.py — CuTile → cubin via tileiras; sm_100+ only.
• _cutedsl.py — CuTeDSL → PTX/cubin via nvidia-cutlass-dsl.

Supporting modules:

• _specs.py — TritonSpec, CuTileSpec, CuTeDSLSpec, KernelImplSpec frozen dataclasses; triton() / cutile() / cutedsl() factories; normalize_impl_to_spec().
• _layer_metadata.py — set_tta_layer_metadata() helper for stamping TRT layer metadata; encode/decode round-trip for custom plugin attribution.
• _recorders.py — launch-parameter recording for Triton/CuTile/CuTeDSL AOT backends.
• _validation.py — spec and descriptor validation utilities.
• _errors.py — TTADiagnosticError structured error type.

Tests (tests/py/annotation/unit/, CPU-only, 46 tests):

• test_specs.py — kernel spec construction, validation, cache-key stability.
• test_specs_custom_plugin.py — CustomPluginSpec and custom_plugin() factory.
• test_layer_metadata.py — metadata encode/decode round-trip.

Design notes

• Descriptor-only, zero core impact. CustomPluginSpec is a plain frozen dataclass. No hooks into _compiler.py, _TRTInterpreter.py, or any other existing file. The integration point (passing a descriptor to a converter) is left for a follow-up PR.
• Deterministic op naming. The QDP op_name is derived from a hash of the kernel function identity, config set, and weight count. The same descriptor created in two different processes produces the same name, making engine caching safe.
• Process-global registration with xdist safety. TRT's QDP registry is process-global. A double-checked locking pattern over a process-level set prevents duplicate registration when pytest-xdist workers share a process.
• Multi-tactic autotuning. Multiple configs dicts produce multiple QDP tactics; TRT's autotuner benchmarks all of them at engine-build time.

trt_plugins.custom_op integration — a torch_tensorrt.dynamo.conversion.plugins API that wires a CustomPluginSpec directly to a registered torch.library custom op, so the TRT lowering path is set up with no manual converter code:

import torch_tensorrt.annotation as tta
from torch_tensorrt.dynamo.conversion import plugins as trt_plugins

def my_op_meta(x, y):
    return x.new_empty(x.shape)

def launch_triton(x, y, out, *, BLOCK: int):
    triton_kernel[...](x, y, out, BLOCK=BLOCK)

def launch_cutile(x, y, out, *, TILE: int):
    cutile_kernel.run(x, y, out, TILE=TILE)

triton_spec = tta.triton(launch_fn=launch_triton, configs=[{"BLOCK": 128}, {"BLOCK": 256}])
cutile_spec = tta.cutile(launch_fn=launch_cutile, configs=[{"TILE": 64}, {"TILE": 128}])

MY_IMPL = tta.custom_plugin(
    kernel=[triton_spec, cutile_spec],
    meta_impl=my_op_meta,
)

trt_plugins.custom_op(
    "torchtrt_ex::my_custom_op",
    impl=MY_IMPL,
)

TRT's autotuner benchmarks all tactics across both backends and selects the fastest for the target GPU.

Future work

This PR establishes the descriptor and registration layer. The follow-up work:

1. Converter integration — wire CustomPluginSpec into the _TRTInterpreter converter dispatch so that annotated subgraphs are lowered to the registered QDP op during torch_tensorrt.compile.
2. Region annotation API — a tta.lower_as(impl=..., name=...) context manager that tags subgraph regions during torch.export for targeted lowering to custom plugins. The intended end-to-end usage looks like:

import torch
import torch_tensorrt
import torch_tensorrt.annotation as tta

@triton.jit
def _fused_add_relu(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = i < n
    tl.store(out_ptr + i, tl.maximum(0, tl.load(x_ptr+i, mask=mask) + tl.load(y_ptr+i, mask=mask)), mask=mask)

def launch_fused_add_relu(x, y, out, stream, BLOCK=256):
    _fused_add_relu[(triton.cdiv(x.numel(), BLOCK),)](x, y, out, x.numel(), BLOCK=BLOCK)

fused_add_relu = tta.custom_plugin(tta.triton(launch_fused_add_relu, configs=[{"BLOCK": 128}, {"BLOCK": 256}]))

class ResidualBlock(nn.Module):
    def forward(self, x, y):
        with tta.lower_as(impl=fused_add_relu):
            return torch.relu(x + y)

model = ResidualBlock().cuda().eval()
x = torch.randn(4, 1024, device="cuda")
trt_model = torch_tensorrt.compile(model, inputs=[x, x])

3. End-to-end tests — GPU tests exercising the full compile → run path for Triton, CuTile, and CuTeDSL backends on representative kernels (fused add+ReLU, RMSNorm, attention).

Test plan

docker exec torch_tensorrt_dev \
  python -m pytest tests/py/annotation/unit/ -n 4 --tb=short -v
# 46 passed